package com.yth.utils;

import org.graalvm.polyglot.Context;
import org.graalvm.polyglot.EnvironmentAccess;
import org.graalvm.polyglot.HostAccess;
import org.graalvm.polyglot.Value;
import org.graalvm.polyglot.proxy.ProxyExecutable;

import java.math.BigDecimal;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Function;

/**
 * Graalvm EJS模板引擎
 *
 * @author yutianhong
 * @version 1.0
 * @since 2025/7/2 14:42
 */
@SuppressWarnings("StringBufferReplaceableByString")
public class GraalvmEJSEngine implements AutoCloseable {
    private final Function<Object, String> toJsonStringFunction;
    private final Context jsContext;
    private final Map<String, Object> registeredFunctions;

    /**
     * @param ejsLibrary           ejs.min.js
     * @param toJsonStringFunction 将Java对象转换为JSON字符串
     */
    public GraalvmEJSEngine(String ejsLibrary, Function<Object, String> toJsonStringFunction) {
        this.toJsonStringFunction = toJsonStringFunction;
        this.registeredFunctions = new ConcurrentHashMap<>();
        this.jsContext = initializeEngine(ejsLibrary);
    }

    private Context initializeEngine(String ejsLibrary) {
        Context jsContext = null;
        try {
            // 创建Graalvm上下文
            jsContext = Context.newBuilder("js")
                    .allowAllAccess(false)
                    .allowIO(false)             // 禁止文件I/O
                    .allowCreateThread(false)   // 禁止创建线程
                    .allowCreateProcess(false)  // 禁止创建进程
                    .allowEnvironmentAccess(EnvironmentAccess.NONE)     // 禁止访问环境变量
                    .allowNativeAccess(false)        //禁止调用本地代码
                    // 只允许明确指定的主机访问
                    .allowHostAccess(HostAccess.newBuilder()
                            .denyAccess(System.class)
                            .denyAccess(Runtime.class)
                            .denyAccess(ProcessBuilder.class)
                            .allowPublicAccess(true)     // 允许访问public方法
                            .allowArrayAccess(true)      // 允许数组访问
                            .allowListAccess(true)       // 允许List访问
                            .allowMapAccess(true)        // 允许Map访问
                            .build())
                    // 设置资源限制
                    .option("engine.WarnInterpreterOnly", "false")
                    .build();

            // 加载EJS库
            jsContext.eval("js", ejsLibrary);

            // 加载辅助函数
            String helperFunctions = loadHelperFunctions();
            jsContext.eval("js", helperFunctions);
            return jsContext;
        } catch (Throwable e) {
            if (jsContext != null) {
                jsContext.close();
            }
            throw new RuntimeException("Failed to initialize Graalvm EJS engine", e);
        }
    }

    /**
     * 注册Java函数到JavaScript上下文中
     *
     * @param functionName JavaScript中的函数名
     * @param function     Java函数
     */
    public void registerFunction(String functionName, Function<Object[], Object> function) {
        Object previous = registeredFunctions.putIfAbsent(functionName, function);
        if (previous != null) {
            throw new IllegalArgumentException("Function " + functionName + " already registered");
        }
        jsContext.getBindings("js").putMember(functionName, (ProxyExecutable) arguments -> {
            Object[] args = new Object[arguments.length];
            for (int i = 0; i < arguments.length; i++) {
                args[i] = parseDynamicValue(arguments[i]); // 将 Polyglot value 转换为 Java 对象
            }
            return function.apply(args); // 执行 Java 函数逻辑，并返回结果
        });
    }

    /**
     * 移除注册的函数
     */
    public void unregisterFunction(String functionName) {
        Object remove = registeredFunctions.remove(functionName);
        if (remove != null) {
            jsContext.getBindings("js").removeMember(functionName);
        }
    }

    /**
     * 获取注册的函数
     */
    public Set<String> registeredFunctions() {
        return registeredFunctions.keySet();
    }

    /**
     * 渲染EJS模板
     */
    public String render(String template, Object data) {
        return render(template, data, null);
    }

    /**
     * 渲染EJS模板（带选项）
     */
    public String render(String template, Object data, Map<String, Object> options) {
        validateTemplate(template);
        try {
            // 将Java对象转换为JSON字符串，然后在JS中解析
            String jsonData = toJsonStringFunction.apply(data);
            String jsonOptions = options != null ? toJsonStringFunction.apply(options) : "{}";

            // 调用JavaScript函数
            StringBuilder renderCode = new StringBuilder();
            renderCode.append("(function() {\n");
            renderCode.append("    var data = JSON.parse(\"").append(escapeForJS(jsonData)).append("\");\n");
            renderCode.append("    var options = JSON.parse(\"").append(escapeForJS(jsonOptions)).append("\");\n");
            renderCode.append("    return safeRender(\"").append(escapeForJS(template)).append("\", data, options);\n");
            renderCode.append("})()");

            Value result = jsContext.eval("js", renderCode.toString());

            // 处理结果
            if (result.getMember("success").asBoolean()) {
                return result.getMember("result").asString();
            } else {
                String error = result.getMember("error").asString();
                throw new EJSRenderException("EJS render failed: " + error);
            }
        } catch (Exception e) {
            throw new EJSRenderException("Failed to render template", e);
        }
    }

    /**
     * 预编译模板
     */
    public CompiledTemplate compile(String template) {
        validateTemplate(template);
        try {
            StringBuilder compileCode = new StringBuilder();
            compileCode.append("(function() {\n");
            compileCode.append("    return compileTemplate(\"").append(escapeForJS(template)).append("\", {});\n");
            compileCode.append("})()");

            Value result = jsContext.eval("js", compileCode.toString());

            if (result.getMember("success").asBoolean()) {
                Value compiled = result.getMember("compiled");
                return new CompiledTemplate(compiled, this);
            } else {
                String error = result.getMember("error").asString();
                throw new EJSRenderException("EJS compile failed: " + error);
            }
        } catch (Exception e) {
            throw new EJSRenderException("Failed to compile template", e);
        }
    }

    @Override
    public void close() {
        if (jsContext != null) {
            jsContext.close();
        }
    }

    private void validateTemplate(String template) {
        if (template == null) {
            throw new IllegalArgumentException("Template cannot be null");
        }
    }

    private String loadHelperFunctions() {
        StringBuilder helperFunctions = new StringBuilder();
        helperFunctions.append("function safeRender(template, data, options) {\n");
        helperFunctions.append("    try {\n");
        helperFunctions.append("        var opts = options || {};\n");
        helperFunctions.append("        opts.async = false;\n");
        helperFunctions.append("        return {\n");
        helperFunctions.append("            success: true,\n");
        helperFunctions.append("            result: ejs.render(template, data, opts)\n");
        helperFunctions.append("        };\n");
        helperFunctions.append("    } catch (error) {\n");
        helperFunctions.append("        return {\n");
        helperFunctions.append("            success: false,\n");
        helperFunctions.append("            error: error.message,\n");
        helperFunctions.append("            stack: error.stack\n");
        helperFunctions.append("        };\n");
        helperFunctions.append("    }\n");
        helperFunctions.append("}\n");

        helperFunctions.append("function compileTemplate(template, options) {\n");
        helperFunctions.append("    try {\n");
        helperFunctions.append("        var opts = options || {};\n");
        helperFunctions.append("        var compiled = ejs.compile(template, opts);\n");
        helperFunctions.append("        return {\n");
        helperFunctions.append("            success: true,\n");
        helperFunctions.append("            compiled: compiled\n");
        helperFunctions.append("        };\n");
        helperFunctions.append("    } catch (error) {\n");
        helperFunctions.append("return {\n");
        helperFunctions.append("            success: false,\n");
        helperFunctions.append("error: error.message\n");
        helperFunctions.append("        };\n");
        helperFunctions.append("    }\n");
        helperFunctions.append("}");
        return helperFunctions.toString();
    }

    private String escapeForJS(String str) {
        String apply = toJsonStringFunction.apply(str);
        return apply.substring(1, apply.length() - 1);
    }

    private Object parseDynamicValue(Value value) {
        if (value.isBoolean()) {
            // 如果是布尔值
            return value.asBoolean();
        } else if (value.isNumber()) {
            // 如果是数字，处理可能的小数
            if (value.fitsInInt()) {
                return value.asInt();  // 如果可以转换为 int
            } else if (value.fitsInLong()) {
                return value.asLong();  // 如果可以转换为 long
            } else if (value.fitsInDouble()) {
                return value.asDouble();  // 如果可以转换为 double
            } else {
                return new BigDecimal(value.asString());
            }
        } else if (value.isString()) {
            // 如果是字符串
            return value.asString();
        } else if (value.isNull()) {
            // 如果是 null
            return null;
        } else if (value.hasArrayElements()) {
            // 如果是数组，使用递归解析每个元素
            List<Object> list = new ArrayList<>();
            for (int i = 0; i < value.getArraySize(); i++) {
                list.add(parseDynamicValue(value.getArrayElement(i)));
            }
            return list;
        } else if (value.hasMembers()) {
            // 如果是对象，使用递归解析每个属性
            Map<String, Object> map = new HashMap<>();
            for (String key : value.getMemberKeys()) {
                map.put(key, parseDynamicValue(value.getMember(key)));
            }
            return map;
        } else if (value.isHostObject()) {
            // 如果是主机语言对象，直接返回它
            return value.asHostObject();
        } else if (value.canExecute()) {
            // 如果是函数
            return (Function<Object[], Object>) value::execute;
        }
        // 如果是未知类型，返回一个通用描述
        return "UnknownValue<" + value + ">";
    }

    // 内部类：编译后的模板
    public static class CompiledTemplate {
        private final Value compiledFunction;
        private final GraalvmEJSEngine engine;

        CompiledTemplate(Value compiledFunction, GraalvmEJSEngine engine) {
            this.compiledFunction = compiledFunction;
            this.engine = engine;
        }

        public String render(Object data) {
            try {
                String jsonData = engine.toJsonStringFunction.apply(data);
                StringBuilder executeCode = new StringBuilder();
                executeCode.append("(function() {\n");
                executeCode.append("    var data = JSON.parse(\"").append(engine.escapeForJS(jsonData)).append("\");\n");
                executeCode.append("    try {\n");
                executeCode.append("        return { success: true, result: arguments[0](data) };\n");
                executeCode.append("    } catch (error) {\n");
                executeCode.append("        return { success: false, error: error.message };\n");
                executeCode.append("    }\n");
                executeCode.append("})");
                Value result = engine.jsContext.eval("js", executeCode.toString()).execute(compiledFunction);

                if (result.getMember("success").asBoolean()) {
                    return result.getMember("result").asString();
                } else {
                    String error = result.getMember("error").asString();
                    throw new EJSRenderException("Compiled template render failed: " + error);
                }
            } catch (Exception e) {
                throw new EJSRenderException("Failed to render compiled template", e);
            }
        }
    }

    // 自定义异常类
    public static class EJSRenderException extends RuntimeException {
        public EJSRenderException(String message) {
            super(message);
        }

        public EJSRenderException(String message, Throwable cause) {
            super(message, cause);
        }
    }
}
